# @Time : 2021/8/4 17:54
# @Author : Li Kunlun
# @Description : GPU计算

import mxnet as mx
from mxnet import nd
from mxnet.gluon import nn

# cpu(0) gpu(0) gpu(1)
print(mx.cpu(), mx.gpu(), mx.gpu(1))

# 1、NDArray的GPU计算
x = nd.array([1, 2, 3])
# 通过NDArray的context属性来查看该NDArray所在的设备
# 输出：cpu(0)
print(x.context)

# 2、GPU上的存储
a = nd.array([1, 2, 3], ctx=mx.gpu())  # 1、创建时指定
# [ 1.  2.  3.]
# <NDArray 3 @gpu(0)>
print(a)

"""
1、通过copyto函数和as_in_context函数在设备之间传输数据,
将内存上的NDArray变量x复制到gpu(0)上。
2、在刚开始时，gpu运算卡顿，需要等待1-2分钟
3、如果源变量和目标变量的context一致，as_in_context函数使目标变量和源变量共享源变量的内存或显存；
    copyto函数总是为目标变量开新的内存或显存。
"""
y = x.copyto(mx.gpu())
z = x.as_in_context(mx.gpu())
# y->
# [ 1.  2.  3.]
# <NDArray 3 @gpu(0)> z ->
# [ 1.  2.  3.]
# <NDArray 3 @gpu(0)>
print("y-> ", y, "z -> ", z)

# False
# True
print(y.copyto(mx.gpu()) is y)
print(y.as_in_context(mx.gpu()) is y)

# 3、GPU上的计算
# [  20.08553696  109.19629669  445.23950195]
# <NDArray 3 @gpu(0)>
print((z + 2).exp() * y)

# 4、Gluon的GPU计算
# 将模型参数初始化在显存上
net = nn.Sequential()
net.add(nn.Dense(1))
net.initialize(ctx=mx.gpu())

# print(net(y)) # print输出会报错，原因未知
net(y)

# [[ 0.0068339]]
# <NDArray 1x1 @gpu(0)>
print(net[0].weight.data())

